package stream

import "sort"

type Row []string

type Stream struct {
	rows []Row
}

type FilterFunc func(Row) bool
type MapFunc func(Row) Row
type HandleFunc func(Row) error

func New(paths ...string) (*Stream, error) {
	s := &Stream{}
	list, err := ReadCSV(paths...)
	if err != nil {
		return nil, err
	}
	s.rows = list
	return s, nil
}

func (s *Stream) Filter(f FilterFunc) *Stream {
	var list []Row
	for _, v := range s.rows {
		if f(v) {
			list = append(list, v)
		}
	}
	s.rows = list
	return s
}

func (s *Stream) Skip(f FilterFunc) *Stream {
	var list []Row
	for _, v := range s.rows {
		if !f(v) {
			list = append(list, v)
		}
	}
	s.rows = list
	return s
}

func (s *Stream) Map(f MapFunc) *Stream {
	var list []Row
	for _, v := range s.rows {
		list = append(list, f(v))
	}
	s.rows = list
	return s
}

func (s *Stream) Distinct(index ...int) *Stream {
	var set = make(map[string]bool)
	var list []Row
	for _, v := range s.rows {
		key := ""
		for _, i := range index {
			key = key + "/" + v[i]
		}
		if set[key] {
			continue
		}
		list = append(list, v)
		set[key] = true
	}
	s.rows = list
	return s
}

func (s *Stream) Sort(f CompareFunc) *Stream {
	sort.Sort(sortBase{s.rows, f})
	return s
}

type sortBase struct {
	rows    []Row
	compare CompareFunc
}

type CompareFunc func(left Row, right Row) bool

func (s sortBase) Len() int           { return len(s.rows) }
func (s sortBase) Less(i, j int) bool { return s.compare(s.rows[i], s.rows[j]) }
func (s sortBase) Swap(i, j int)      { s.rows[i], s.rows[j] = s.rows[j], s.rows[i] }

func (s *Stream) Collect() []Row {
	return s.rows
}

func (s *Stream) Values(index int) Row {
	var values Row
	for _, v := range s.rows {
		values = append(values, v[index])
	}
	return values
}

func (s *Stream) Handle(f HandleFunc) error {
	for _, v := range s.rows {
		err := f(v)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Stream) Save(path string, header ...Row) error {
	return WriteCSV(path, s.rows, header...)
}
